{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T04:22:38.285207Z",
     "start_time": "2025-07-31T04:22:38.279490Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "from scipy.linalg import lstsq\n",
    "path = \"D:/7_30data/\""
   ],
   "id": "435a83cfe4efaee4",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T04:22:38.324958Z",
     "start_time": "2025-07-31T04:22:38.310560Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def iterative_refinement_solve(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True):\n",
    "    \"\"\"\n",
    "    迭代精化求解复数线性系统 S @ a = b，其中 S = V^H @ V, b = V^H @ L_rho\n",
    "\n",
    "    Parameters:\n",
    "    -----------\n",
    "    V : ndarray, shape (m, n)\n",
    "        输入复数矩阵\n",
    "    L_rho : ndarray, shape (m,)\n",
    "        右端复数向量\n",
    "    max_iter : int\n",
    "        最大迭代次数\n",
    "    tolerance : float\n",
    "        收敛容差\n",
    "    verbose : bool\n",
    "        是否输出详细信息\n",
    "\n",
    "    Returns:\n",
    "    --------\n",
    "    a : ndarray\n",
    "        解向量 (复数形式)\n",
    "    \"\"\"\n",
    "\n",
    "    # 构建正规方程 S @ a = b\n",
    "    S = V.conj().T @ V  # V^H @ V\n",
    "    b = V.conj().T @ L_rho  # V^H @ L_rho\n",
    "\n",
    "    # 计算并输出矩阵的秩信息\n",
    "    rank_S = np.linalg.matrix_rank(S)\n",
    "    augmented = np.hstack((S, b.reshape(-1, 1)))\n",
    "    rank_augmented = np.linalg.matrix_rank(augmented)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"rank(S) = {rank_S}\")\n",
    "        print(f\"rank([S | b]) = {rank_augmented}\")\n",
    "\n",
    "        # 输出条件数信息\n",
    "        cond_S = np.linalg.cond(S)\n",
    "        print(f\"Condition number of S: {cond_S:.4e}\")\n",
    "        print(\"\\n=== Iterative Refinement Process ===\")\n",
    "\n",
    "    # 检查系统的相容性\n",
    "    if rank_S != rank_augmented:\n",
    "        print(\"Warning: System is inconsistent (rank(S) ≠ rank([S|b]))\")\n",
    "        print(\"Using least squares solution...\")\n",
    "\n",
    "    # 步骤1: 获得初始解 (使用最小二乘法)\n",
    "    try:\n",
    "        a = lstsq(S, b, lapack_driver='gelsy')[0]\n",
    "        if verbose:\n",
    "            print(\"Initial solution computed using least squares\")\n",
    "    except Exception as e:\n",
    "        print(f\"Failed to compute initial solution: {e}\")\n",
    "        return np.zeros(S.shape[1], dtype=complex)\n",
    "\n",
    "    # 初始残差\n",
    "    initial_residual = np.linalg.norm(S @ a - b)\n",
    "    if verbose:\n",
    "        print(f\"Initial residual norm: {initial_residual:.4e}\")\n",
    "\n",
    "    # 迭代精化过程\n",
    "    for i in range(max_iter):\n",
    "        # 步骤2: 计算当前残差 r = b - S @ a\n",
    "        r = b - S @ a\n",
    "        residual_norm = np.linalg.norm(r)\n",
    "\n",
    "        if verbose:\n",
    "            print(f\"Iteration {i+1}: Residual norm = {residual_norm:.4e}\")\n",
    "\n",
    "        # 步骤3: 检查收敛\n",
    "        if residual_norm < tolerance:\n",
    "            if verbose:\n",
    "                print(f\"Converged after {i+1} iterations\")\n",
    "            break\n",
    "\n",
    "        # 步骤4: 求解修正方程 S @ da = r\n",
    "        try:\n",
    "            da = lstsq(S, r, lapack_driver='gelsy')[0]\n",
    "        except Exception as e:\n",
    "            if verbose:\n",
    "                print(f\"Iteration {i+1} failed: {e}\")\n",
    "            break\n",
    "\n",
    "        # 步骤5: 更新解 a = a + da\n",
    "        a = a + da\n",
    "\n",
    "        # 检查是否有改进，如果连续几次迭代没有改进则停止\n",
    "        if i > 2 and residual_norm >= tolerance:\n",
    "            prev_residuals = []\n",
    "            for j in range(max(0, i-2), i+1):\n",
    "                prev_residuals.append(residual_norm)\n",
    "\n",
    "            # 如果最近3次迭代残差变化很小，可能已经达到数值极限\n",
    "            if len(prev_residuals) >= 3:\n",
    "                residual_changes = [abs(prev_residuals[k] - prev_residuals[k-1])\n",
    "                                  for k in range(1, len(prev_residuals))]\n",
    "                if all(change < tolerance * 10 for change in residual_changes):\n",
    "                    if verbose:\n",
    "                        print(f\"Convergence stalled after {i+1} iterations\")\n",
    "                    break\n",
    "\n",
    "    # 计算最终残差\n",
    "    final_residual = np.linalg.norm(S @ a - b)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Final residual norm for S @ a - b: {final_residual:.4e}\")\n",
    "\n",
    "        # 输出改进情况\n",
    "        improvement_factor = initial_residual / final_residual if final_residual > 0 else float('inf')\n",
    "        print(f\"Improvement factor: {improvement_factor:.2e}\")\n",
    "\n",
    "        # 检查解的性质\n",
    "        max_imag = np.max(np.abs(np.imag(a)))\n",
    "        if max_imag > 1e-12:\n",
    "            print(f\"Warning: Solution has significant imaginary parts: {max_imag:.4e}\")\n",
    "\n",
    "    # 输出必要的信息\n",
    "    print(f\"rank(S) = {rank_S}\")\n",
    "    print(f\"rank([S | b]) = {rank_augmented}\")\n",
    "    print(f\"Residual norm for S @ a - b: {final_residual:.4e}\")\n",
    "\n",
    "    return a"
   ],
   "id": "4f853ab2adfc0780",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T04:22:38.351257Z",
     "start_time": "2025-07-31T04:22:38.345085Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def iterative_refinement_solve_real_output(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True):\n",
    "    \"\"\"\n",
    "    迭代精化求解，返回实数解 (如果虚部可忽略)\n",
    "\n",
    "    Parameters:\n",
    "    -----------\n",
    "    V : ndarray, shape (m, n)\n",
    "        输入复数矩阵\n",
    "    L_rho : ndarray, shape (m,)\n",
    "        右端复数向量\n",
    "    max_iter : int\n",
    "        最大迭代次数\n",
    "    tolerance : float\n",
    "        收敛容差\n",
    "    verbose : bool\n",
    "        是否输出详细信息\n",
    "\n",
    "    Returns:\n",
    "    --------\n",
    "    a : ndarray\n",
    "        解向量 (如果虚部可忽略则返回实数，否则返回复数)\n",
    "    \"\"\"\n",
    "\n",
    "    # 调用复数版本\n",
    "    a_complex = iterative_refinement_solve(V, L_rho, max_iter, tolerance, verbose)\n",
    "\n",
    "    # 检查是否可以转换为实数\n",
    "    max_imag = np.max(np.abs(np.imag(a_complex)))\n",
    "\n",
    "    if max_imag < 1e-12:  # 虚部可忽略\n",
    "        if verbose:\n",
    "            print(\"Converting to real solution (imaginary parts negligible)\")\n",
    "        return np.real(a_complex)\n",
    "    else:\n",
    "        if verbose:\n",
    "            print(f\"Keeping complex solution (max imaginary part: {max_imag:.4e})\")\n",
    "        return a_complex"
   ],
   "id": "3ce980e7ee9caa8d",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T04:22:38.380149Z",
     "start_time": "2025-07-31T04:22:38.375903Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 简化版本 - 直接替换您原来的函数\n",
    "def improved_iterative_solve(V, L_rho):\n",
    "    \"\"\"\n",
    "    简化版迭代精化求解，直接替换您的原函数\n",
    "    输出格式与您的需求完全一致\n",
    "    \"\"\"\n",
    "    return iterative_refinement_solve_real_output(V, L_rho, max_iter=15, tolerance=1e-15, verbose=True)"
   ],
   "id": "f4f26ba57c2e99ae",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T04:23:37.548881Z",
     "start_time": "2025-07-31T04:23:35.951524Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 使用示例\n",
    "V = np.load(path+'V.npy').astype(np.complex128)\n",
    "L_rho = np.load(path+'L_rho.npy').astype(np.complex128)\n",
    "\n",
    "    #调用函数\n",
    "solution = improved_iterative_solve(V, L_rho)\n",
    "\n",
    "    #或者使用详细版本\n",
    "solution = iterative_refinement_solve(V, L_rho, max_iter=20, tolerance=1e-16)\n",
    "\n",
    "pass"
   ],
   "id": "63a9254cdc5c9528",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rank(S) = 2\n",
      "rank([S | b]) = 2\n",
      "Condition number of S: 9.7122e+80\n",
      "\n",
      "=== Iterative Refinement Process ===\n",
      "Initial solution computed using least squares\n",
      "Initial residual norm: 1.4462e-10\n",
      "Iteration 1: Residual norm = 1.4462e-10\n",
      "Iteration 2: Residual norm = 2.0580e-11\n",
      "Iteration 3: Residual norm = 8.2318e-11\n",
      "Iteration 4: Residual norm = 2.0580e-11\n",
      "Convergence stalled after 4 iterations\n",
      "Final residual norm for S @ a - b: 2.0580e-11\n",
      "Improvement factor: 7.03e+00\n",
      "rank(S) = 2\n",
      "rank([S | b]) = 2\n",
      "Residual norm for S @ a - b: 2.0580e-11\n",
      "Converting to real solution (imaginary parts negligible)\n",
      "rank(S) = 2\n",
      "rank([S | b]) = 2\n",
      "Condition number of S: 9.7122e+80\n",
      "\n",
      "=== Iterative Refinement Process ===\n",
      "Initial solution computed using least squares\n",
      "Initial residual norm: 1.4462e-10\n",
      "Iteration 1: Residual norm = 1.4462e-10\n",
      "Iteration 2: Residual norm = 2.0580e-11\n",
      "Iteration 3: Residual norm = 8.2318e-11\n",
      "Iteration 4: Residual norm = 2.0580e-11\n",
      "Convergence stalled after 4 iterations\n",
      "Final residual norm for S @ a - b: 2.0580e-11\n",
      "Improvement factor: 7.03e+00\n",
      "rank(S) = 2\n",
      "rank([S | b]) = 2\n",
      "Residual norm for S @ a - b: 2.0580e-11\n"
     ]
    }
   ],
   "execution_count": 13
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
